Phase Mask Design¶

In this notebook, we will illustrate the problem of inverse design of a phase mask: we will choose the example from Wong et al, 2021, designing a diffractive pupil phase mask for the Toliman telescope.

In order to get high precision centroids, we need to maximize the gradient energy of the pupil; in order to satisfy fabrication constraints, we need a binary mask with phases only in {0, π}.

In [1]:

# Core jax
import jax
import jax.numpy as np
import jax.random as jr

# Optimisation
import equinox as eqx
import optax

# Optics
import dLux as dl

# Plotting/visualisation
import matplotlib.pyplot as plt
from tqdm.notebook import tqdm

%matplotlib inline
plt.rcParams['image.cmap'] = 'inferno'
plt.rcParams["font.family"] = "serif"
plt.rcParams['figure.dpi'] = 120

# Core jax import jax import jax.numpy as np import jax.random as jr # Optimisation import equinox as eqx import optax # Optics import dLux as dl # Plotting/visualisation import matplotlib.pyplot as plt from tqdm.notebook import tqdm %matplotlib inline plt.rcParams['image.cmap'] = 'inferno' plt.rcParams["font.family"] = "serif" plt.rcParams['figure.dpi'] = 120

/Users/louis/mambaforge/envs/dlux/lib/python3.10/site-packages/chex/_src/pytypes.py:37: FutureWarning: jax.tree_structure is deprecated, and will be removed in a future release. Use jax.tree_util.tree_structure instead.
  PyTreeDef = type(jax.tree_structure(None))

We will first generate an orthonormal basis for the pupil phases, and then threshold this to {0, 1} while preserving soft edges using the Continuous Latent Image Mask Binarization (CLIMB) algorithm from the Wong et al paper.

Generate basis vectors however you like - in this case we are using logarithmic radial harmonics and sines and cosines in θ., but you can do whatever you like here. This code is not important; just generate your favourite not-necessarily-orthonormal basis, and we will use PCA to orthonormalize it later on.

In [2]:

# Define arrays sizes, samplings, symmetries
wf_npix = 256
oversample = 3
nslice = 3

# Define basis hyper parameters
a = 10
b = 8
ith = 10

# Define arrays sizes, samplings, symmetries wf_npix = 256 oversample = 3 nslice = 3 # Define basis hyper parameters a = 10 b = 8 ith = 10

In [3]:

# Define coordinate grids
npix = wf_npix * oversample
c = (npix - 1) / 2.
xs = (np.arange(npix) - c) / c
XX, YY = np.meshgrid(xs, xs)
RR = np.sqrt(XX ** 2 + YY ** 2)
PHI = np.arctan2(YY, XX)

# Generate basis vectors to map over
As = np.arange(-a, a+1)
Bs = nslice * np.arange(0, b+1)
Cs = np.array([-np.pi/2, np.pi/2])
Is = np.arange(-ith, ith+1)

# Define basis functions
LRHF_fn = lambda A, B, C, RR, PHI: np.cos(A*np.log(RR + 1e-12) + B*PHI + C)
sine_fn = lambda i, RR: np.sin(i * np.pi * RR)
cose_fn = lambda i, RR: np.cos(i * np.pi * RR)

# Map over basis functions
gen_LRHF_basis = jax.vmap(jax.vmap(jax.vmap( \
                            LRHF_fn, (None, 0, None, None, None)), 
                                     (0, None, None, None, None)), 
                                     (None, None, 0, None, None))
gen_sine_basis = jax.vmap(sine_fn, in_axes=(0, None))
gen_cose_basis = jax.vmap(cose_fn, in_axes=(0, None))

# Generate basis
LRHF_basis = gen_LRHF_basis(As, Bs, Cs, RR, PHI) \
            .reshape([len(As)*len(Bs)*len(Cs), npix, npix])
sine_basis = gen_sine_basis(Is, RR)
cose_basis = gen_cose_basis(Is, RR)

# Format shapes and combine
LRHF_flat = LRHF_basis.reshape([len(As)*len(Bs)*len(Cs), npix*npix])
sine_flat = sine_basis.reshape([len(sine_basis), npix*npix])
cose_flat = cose_basis.reshape([len(cose_basis), npix*npix])
full_basis = np.concatenate([
    LRHF_flat,
    sine_flat,
    cose_flat
])

# Define coordinate grids npix = wf_npix * oversample c = (npix - 1) / 2. xs = (np.arange(npix) - c) / c XX, YY = np.meshgrid(xs, xs) RR = np.sqrt(XX ** 2 + YY ** 2) PHI = np.arctan2(YY, XX) # Generate basis vectors to map over As = np.arange(-a, a+1) Bs = nslice * np.arange(0, b+1) Cs = np.array([-np.pi/2, np.pi/2]) Is = np.arange(-ith, ith+1) # Define basis functions LRHF_fn = lambda A, B, C, RR, PHI: np.cos(A*np.log(RR + 1e-12) + B*PHI + C) sine_fn = lambda i, RR: np.sin(i * np.pi * RR) cose_fn = lambda i, RR: np.cos(i * np.pi * RR) # Map over basis functions gen_LRHF_basis = jax.vmap(jax.vmap(jax.vmap( \ LRHF_fn, (None, 0, None, None, None)), (0, None, None, None, None)), (None, None, 0, None, None)) gen_sine_basis = jax.vmap(sine_fn, in_axes=(0, None)) gen_cose_basis = jax.vmap(cose_fn, in_axes=(0, None)) # Generate basis LRHF_basis = gen_LRHF_basis(As, Bs, Cs, RR, PHI) \ .reshape([len(As)*len(Bs)*len(Cs), npix, npix]) sine_basis = gen_sine_basis(Is, RR) cose_basis = gen_cose_basis(Is, RR) # Format shapes and combine LRHF_flat = LRHF_basis.reshape([len(As)*len(Bs)*len(Cs), npix*npix]) sine_flat = sine_basis.reshape([len(sine_basis), npix*npix]) cose_flat = cose_basis.reshape([len(cose_basis), npix*npix]) full_basis = np.concatenate([ LRHF_flat, sine_flat, cose_flat ])

Orthonormalize with PCA - could also use Gram-Schmidt if you prefer.

In [4]:

%%time
from sklearn.decomposition import PCA
pca = PCA().fit(full_basis)

components = pca.components_.reshape([len(full_basis), npix, npix])
components = np.copy(components[:99,:,:])
basis = np.concatenate([np.mean(pca.mean_)*np.array(np.ones((1,npix,npix))), components])

%%time from sklearn.decomposition import PCA pca = PCA().fit(full_basis) components = pca.components_.reshape([len(full_basis), npix, npix]) components = np.copy(components[:99,:,:]) basis = np.concatenate([np.mean(pca.mean_)*np.array(np.ones((1,npix,npix))), components])

CPU times: user 2min 58s, sys: 1.78 s, total: 3min
Wall time: 33.1 s

Show the pretty basis vectors:

In [5]:

nfigs = 100
ncols = 10
nrows = 1 + nfigs//ncols

plt.figure(figsize=(4*ncols, 4*nrows))
for i in range(nfigs):
    plt.subplot(nrows, ncols, i+1)
    plt.imshow(basis[i], cmap='seismic')
    plt.xticks([])
    plt.yticks([])
plt.tight_layout()
plt.show()

nfigs = 100 ncols = 10 nrows = 1 + nfigs//ncols plt.figure(figsize=(4*ncols, 4*nrows)) for i in range(nfigs): plt.subplot(nrows, ncols, i+1) plt.imshow(basis[i], cmap='seismic') plt.xticks([]) plt.yticks([]) plt.tight_layout() plt.show()

In [6]:

# save for use later
np.save('files/basis', basis)

# save for use later np.save('files/basis', basis)

---

Optimizing the Pupil¶

Initialize a wavelength grid and decide on the mode coordinate grid scaling relative to the pupil:

In [7]:

# Define our observation wavelengths
wavels = 1e-9 * np.linspace(595, 695, 3) # wavelengths always in m

# Convert fringe sizes to pixel values
ppf = dl.utils.get_ppf(wavels)
rmin = 1/ppf
rmax = 10/ppf
print(rmin, rmax)

# Define our observation wavelengths wavels = 1e-9 * np.linspace(595, 695, 3) # wavelengths always in m # Convert fringe sizes to pixel values ppf = dl.utils.get_ppf(wavels) rmin = 1/ppf rmax = 10/ppf print(rmin, rmax)

5.2393847 52.39385

Now we initialize an optical system using the default Toliman model, plus the CLIMB layer that soft-thresholds the edges (see Wong et al, 2021, sec 3.2.2). In brief, what this does is make an OPD map as a weighted sum of modes; where this is positive, we set to π phase, where it is negative, we set it to zero, and on the edges, we soft-edge it to propagate gradients.

In [8]:

# We want to the use 'ApplyBasisCLIMB' layer, which lets us optimise a 
# mask comprised of binary values from a continuous basis
coeffs = 100*jr.normal(jr.PRNGKey(0), [len(basis)])
extra_layers = [dl.ApplyBasisCLIMB(basis, coeffs, np.mean(wavels))]

# Gernerate optical system with the extra mask layer
layers = dl.utils.toliman_layers(extra_layers)    
osys = dl.OpticalSystem(layers, wavels=wavels, detector_layers=[dl.ApplyJitter(1)])

# We want to the use 'ApplyBasisCLIMB' layer, which lets us optimise a # mask comprised of binary values from a continuous basis coeffs = 100*jr.normal(jr.PRNGKey(0), [len(basis)]) extra_layers = [dl.ApplyBasisCLIMB(basis, coeffs, np.mean(wavels))] # Gernerate optical system with the extra mask layer layers = dl.utils.toliman_layers(extra_layers) osys = dl.OpticalSystem(layers, wavels=wavels, detector_layers=[dl.ApplyJitter(1)])

Visualize this:

In [9]:

fig, (ax1,ax2) = plt.subplots(1,2,figsize=(6.0,3.3))
ax1.imshow(osys.layers[2].array*osys.layers[4].get_binary_phase(),cmap='hot')
ax1.set_title('Input OPD')

psf = osys.propagate()
ax2.imshow(psf)
ax2.set_title('Input PSF')
plt.show()

fig, (ax1,ax2) = plt.subplots(1,2,figsize=(6.0,3.3)) ax1.imshow(osys.layers[2].array*osys.layers[4].get_binary_phase(),cmap='hot') ax1.set_title('Input OPD') psf = osys.propagate() ax2.imshow(psf) ax2.set_title('Input PSF') plt.show()

Now lets define our loss function. We can use the .get_filter_spec() function in conjunction with the path to our coefficients in order to automatically generate a filter that says we only want to get gradients for those terms!

In [10]:

# Define path to coefficients and get filter spec
coeffs_path = ['layers', 4, 'coeffs']
filter_spec = osys.get_filter_spec([coeffs_path])

@eqx.filter_jit
@eqx.filter_value_and_grad(arg=filter_spec)
def loss_func(tel, rmin=20, rmax=150, power=0.5):
    # Get PSF, Gradients and Mask
    psf = tel.propagate()
    Rmask = dl.utils.get_Rmask(psf.shape[0], rmin, rmax)

    # Calcualte loss
    loss1 = - np.power(Rmask*dl.utils.get_GE(psf), power).sum()
    loss2 = - np.power(Rmask*dl.utils.get_RGE(psf), power).sum()
    loss = loss1 + loss2
    return loss

# Define path to coefficients and get filter spec coeffs_path = ['layers', 4, 'coeffs'] filter_spec = osys.get_filter_spec([coeffs_path]) @eqx.filter_jit @eqx.filter_value_and_grad(arg=filter_spec) def loss_func(tel, rmin=20, rmax=150, power=0.5): # Get PSF, Gradients and Mask psf = tel.propagate() Rmask = dl.utils.get_Rmask(psf.shape[0], rmin, rmax) # Calcualte loss loss1 = - np.power(Rmask*dl.utils.get_GE(psf), power).sum() loss2 = - np.power(Rmask*dl.utils.get_RGE(psf), power).sum() loss = loss1 + loss2 return loss

Evaluate once to jit compile:

In [11]:

%%time
loss, grads = loss_func(osys, rmin=rmin, rmax=rmax) # Compile

print("Initial Loss: {}".format(loss))

%%time loss, grads = loss_func(osys, rmin=rmin, rmax=rmax) # Compile print("Initial Loss: {}".format(loss))

Initial Loss: -34.64957046508789
CPU times: user 4.1 s, sys: 110 ms, total: 4.21 s
Wall time: 2.93 s

Gradient descent time!

In [12]:

model = osys
optim = optax.adam(2.5e1)
opt_state = optim.init(model)

errors, grads_out, models_out = [], [], []

with tqdm(range(100),desc='Gradient Descent') as t:
    for i in t: 
        loss, grads = loss_func(model, rmin=rmin, rmax=rmax)
        updates, opt_state = optim.update(grads, opt_state)
        model = eqx.apply_updates(model, updates)
        models_out.append(model)
        errors.append(loss)
        grads_out.append(grads)

        t.set_description("Loss: {:.3f}".format(loss)) # update the progress bar

model = osys optim = optax.adam(2.5e1) opt_state = optim.init(model) errors, grads_out, models_out = [], [], [] with tqdm(range(100),desc='Gradient Descent') as t: for i in t: loss, grads = loss_func(model, rmin=rmin, rmax=rmax) updates, opt_state = optim.update(grads, opt_state) model = eqx.apply_updates(model, updates) models_out.append(model) errors.append(loss) grads_out.append(grads) t.set_description("Loss: {:.3f}".format(loss)) # update the progress bar

Gradient Descent:   0%|          | 0/100 [00:00<?, ?it/s]

Visualize the results:

In [13]:

coeffs_out = np.array([model.get_leaf(coeffs_path) for model in models_out])

mosaic = """
    AABB
    CDEF
    """

fig = plt.figure(constrained_layout=True,figsize=(12, 8))
axes = fig.subplot_mosaic(mosaic)

for ax in ['C','D','E','F']:
    axes[ax].set_xticks([])
    axes[ax].set_yticks([])


axes['A'].plot(np.array(errors))
axes['A'].set_title("Loss")

axes['B'].set_title("Coeffs")
axes['B'].plot(coeffs_out[:], alpha=0.2)

model_in = osys
axes['C'].imshow(model_in.layers[2].array*model_in.layers[4].get_binary_phase(),cmap='hot')
axes['C'].set_title('Input OPD')

psf_in = model_in.propagate()
axes['D'].imshow(psf_in)
axes['D'].set_title('Input PSF')

final = models_out[-1]
axes['E'].imshow(final.layers[2].array*final.layers[4].get_binary_phase(),cmap='hot')
axes['E'].set_title('Final OPD')

psf_out = final.propagate()
axes['F'].imshow(psf_out)
axes['F'].set_title('Final PSF')

coeffs_out = np.array([model.get_leaf(coeffs_path) for model in models_out]) mosaic = """ AABB CDEF """ fig = plt.figure(constrained_layout=True,figsize=(12, 8)) axes = fig.subplot_mosaic(mosaic) for ax in ['C','D','E','F']: axes[ax].set_xticks([]) axes[ax].set_yticks([]) axes['A'].plot(np.array(errors)) axes['A'].set_title("Loss") axes['B'].set_title("Coeffs") axes['B'].plot(coeffs_out[:], alpha=0.2) model_in = osys axes['C'].imshow(model_in.layers[2].array*model_in.layers[4].get_binary_phase(),cmap='hot') axes['C'].set_title('Input OPD') psf_in = model_in.propagate() axes['D'].imshow(psf_in) axes['D'].set_title('Input PSF') final = models_out[-1] axes['E'].imshow(final.layers[2].array*final.layers[4].get_binary_phase(),cmap='hot') axes['E'].set_title('Final OPD') psf_out = final.propagate() axes['F'].imshow(psf_out) axes['F'].set_title('Final PSF')

Out[13]:

Text(0.5, 1.0, 'Final PSF')

In [14]:

# Save mask for use in flatfield_calibration notebook
mask_out = models_out[-1].layers[4].get_binary_phase()
np.save("files/test_mask", mask_out)

# Save mask for use in flatfield_calibration notebook mask_out = models_out[-1].layers[4].get_binary_phase() np.save("files/test_mask", mask_out)

In [ ]: